{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nBuilding a Graph Convolutional Network\n======================================\n**Author**: `Yulun Yao <https://yulunyao.io/>`_,             `Chien-Yu Lin <https://homes.cs.washington.edu/~cyulin/>`_\n\nThis article is an introductory tutorial to build a Graph Convolutional Network (GCN) with Relay.\nIn this tutorial, we will run our GCN on Cora dataset to demonstrate.\nCora dataset is a common benchmark for Graph Neural Networks (GNN) and frameworks that support GNN training and inference.\nWe directly load the dataset from DGL library to do the apples to apples comparison against DGL.\n\nPlease refer to DGL doc for DGL installation at\nhttps://docs.dgl.ai/install/index.html.\n\nPlease refer to PyTorch guide for PyTorch installation at\nhttps://pytorch.org/get-started/locally/.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define GCN in DGL with PyTorch backend\n--------------------------------------\n\nDGL example: https://github.com/dmlc/dgl/tree/master/examples/pytorch/gcn\nThis part reuses the code from the above example.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport dgl\nimport networkx as nx\nfrom dgl.nn.pytorch import GraphConv\n\n\nclass GCN(nn.Module):\n    def __init__(self, g, n_infeat, n_hidden, n_classes, n_layers, activation):\n        super(GCN, self).__init__()\n        self.g = g\n        self.layers = nn.ModuleList()\n        self.layers.append(GraphConv(n_infeat, n_hidden, activation=activation))\n        for i in range(n_layers - 1):\n            self.layers.append(GraphConv(n_hidden, n_hidden, activation=activation))\n        self.layers.append(GraphConv(n_hidden, n_classes))\n\n    def forward(self, features):\n        h = features\n        for i, layer in enumerate(self.layers):\n            # handle api changes for differnt DGL version\n            if dgl.__version__ > \"0.3\":\n                h = layer(self.g, h)\n            else:\n                h = layer(h, self.g)\n        return h"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define the functions to load dataset and evaluate accuracy\n----------------------------------------------------------\nYou may substitute this part with your own dataset, here we load data from DGL\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from dgl.data import load_data\nfrom collections import namedtuple\n\n\ndef load_dataset(dataset=\"cora\"):\n    args = namedtuple(\"args\", [\"dataset\"])\n    data = load_data(args(dataset))\n\n    # Remove self-loops to avoid duplicate passing of a node's feature to itself\n    g = data.graph\n    g.remove_edges_from(nx.selfloop_edges(g))\n    g.add_edges_from(zip(g.nodes, g.nodes))\n\n    return g, data\n\n\ndef evaluate(data, logits):\n    test_mask = data.test_mask  # the test set which isn't included in the training phase\n\n    pred = logits.argmax(axis=1)\n    acc = ((pred == data.labels) * test_mask).sum() / test_mask.sum()\n\n    return acc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load the data and set up model parameters\n-----------------------------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "\"\"\"\nParameters\n----------\ndataset: str\n    Name of dataset. You can choose from ['cora', 'citeseer', 'pubmed'].\n\nnum_layer: int\n    number of hidden layers\n\nnum_hidden: int\n    number of the hidden units in the hidden layer\n\ninfeat_dim: int\n    dimension of the input features\n\nnum_classes: int\n    dimension of model output (Number of classes)\n\"\"\"\ndataset = \"cora\"\n\ng, data = load_dataset(dataset)\n\nnum_layers = 1\nnum_hidden = 16\ninfeat_dim = data.features.shape[1]\nnum_classes = data.num_labels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Set up the DGL-PyTorch model and get the golden results\n-------------------------------------------------------\n\nThe weights are trained with https://github.com/dmlc/dgl/blob/master/examples/pytorch/gcn/train.py\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib.download import download_testdata\nfrom dgl import DGLGraph\n\nfeatures = torch.FloatTensor(data.features)\ndgl_g = DGLGraph(g)\n\ntorch_model = GCN(dgl_g, infeat_dim, num_hidden, num_classes, num_layers, F.relu)\n\n# Download the pretrained weights\nmodel_url = \"https://homes.cs.washington.edu/~cyulin/media/gnn_model/gcn_%s.torch\" % (dataset)\nmodel_path = download_testdata(model_url, \"gcn_%s.pickle\" % (dataset), module=\"gcn_model\")\n\n# Load the weights into the model\ntorch_model.load_state_dict(torch.load(model_path))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Run the DGL model and test for accuracy\n---------------------------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "torch_model.eval()\nwith torch.no_grad():\n    logits_torch = torch_model(features)\nprint(\"Print the first five outputs from DGL-PyTorch execution\\n\", logits_torch[:5])\n\nacc = evaluate(data, logits_torch.numpy())\nprint(\"Test accuracy of DGL results: {:.2%}\".format(acc))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define Graph Convolution Layer in Relay\n---------------------------------------\nTo run GCN on TVM, we first need to implement Graph Convolution Layer.\nYou may refer to https://github.com/dmlc/dgl/blob/master/python/dgl/nn/mxnet/conv/graphconv.py for a GraphConv Layer implemented in DGL with MXNet Backend\n\nThe layer is defined with below operations, note that we apply two transposes to keep adjacency matrix on right hand side of sparse_dense operator,\nthis method is temporary and will be updated in next few weeks when we have sparse matrix transpose and support for left sparse operator.\n\n .. math::\n\n           \\mbox{GraphConv}(A, H, W)   = A * H * W\n                                       = ((H * W)^t * A^t)^t\n                                       = ((W^t * H^t) * A^t)^t\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm import relay\nfrom tvm.contrib import graph_executor\nimport tvm\nfrom tvm import te\n\n\ndef GraphConv(layer_name, input_dim, output_dim, adj, input, norm=None, bias=True, activation=None):\n    \"\"\"\n    Parameters\n    ----------\n    layer_name: str\n    Name of layer\n\n    input_dim: int\n    Input dimension per node feature\n\n    output_dim: int,\n    Output dimension per node feature\n\n    adj: namedtuple,\n    Graph representation (Adjacency Matrix) in Sparse Format (`data`, `indices`, `indptr`),\n    where `data` has shape [num_nonzeros], indices` has shape [num_nonzeros], `indptr` has shape [num_nodes + 1]\n\n    input: relay.Expr,\n    Input feature to current layer with shape [num_nodes, input_dim]\n\n    norm: relay.Expr,\n    Norm passed to this layer to normalize features before and after Convolution.\n\n    bias: bool\n    Set bias to True to add bias when doing GCN layer\n\n    activation: <function relay.op.nn>,\n    Activation function applies to the output. e.g. relay.nn.{relu, sigmoid, log_softmax, softmax, leaky_relu}\n\n    Returns\n    ----------\n    output: tvm.relay.Expr\n    The Output Tensor for this layer [num_nodes, output_dim]\n    \"\"\"\n    if norm is not None:\n        input = relay.multiply(input, norm)\n\n    weight = relay.var(layer_name + \".weight\", shape=(input_dim, output_dim))\n    weight_t = relay.transpose(weight)\n    dense = relay.nn.dense(weight_t, input)\n    output = relay.nn.sparse_dense(dense, adj)\n    output_t = relay.transpose(output)\n    if norm is not None:\n        output_t = relay.multiply(output_t, norm)\n    if bias is True:\n        _bias = relay.var(layer_name + \".bias\", shape=(output_dim, 1))\n        output_t = relay.nn.bias_add(output_t, _bias, axis=-1)\n    if activation is not None:\n        output_t = activation(output_t)\n    return output_t"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Prepare the parameters needed in the GraphConv layers\n-----------------------------------------------------\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport networkx as nx\n\n\ndef prepare_params(g, data):\n    params = {}\n    params[\"infeats\"] = data.features.numpy().astype(\n        \"float32\"\n    )  # Only support float32 as feature for now\n\n    # Generate adjacency matrix\n    adjacency = nx.to_scipy_sparse_matrix(g)\n    params[\"g_data\"] = adjacency.data.astype(\"float32\")\n    params[\"indices\"] = adjacency.indices.astype(\"int32\")\n    params[\"indptr\"] = adjacency.indptr.astype(\"int32\")\n\n    # Normalization w.r.t. node degrees\n    degs = [g.in_degree[i] for i in range(g.number_of_nodes())]\n    params[\"norm\"] = np.power(degs, -0.5).astype(\"float32\")\n    params[\"norm\"] = params[\"norm\"].reshape((params[\"norm\"].shape[0], 1))\n\n    return params\n\n\nparams = prepare_params(g, data)\n\n# Check shape of features and the validity of adjacency matrix\nassert len(params[\"infeats\"].shape) == 2\nassert (\n    params[\"g_data\"] is not None and params[\"indices\"] is not None and params[\"indptr\"] is not None\n)\nassert params[\"infeats\"].shape[0] == params[\"indptr\"].shape[0] - 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Put layers together\n-------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Define input features, norms, adjacency matrix in Relay\ninfeats = relay.var(\"infeats\", shape=data.features.shape)\nnorm = relay.Constant(tvm.nd.array(params[\"norm\"]))\ng_data = relay.Constant(tvm.nd.array(params[\"g_data\"]))\nindices = relay.Constant(tvm.nd.array(params[\"indices\"]))\nindptr = relay.Constant(tvm.nd.array(params[\"indptr\"]))\n\nAdjacency = namedtuple(\"Adjacency\", [\"data\", \"indices\", \"indptr\"])\nadj = Adjacency(g_data, indices, indptr)\n\n# Construct the 2-layer GCN\nlayers = []\nlayers.append(\n    GraphConv(\n        layer_name=\"layers.0\",\n        input_dim=infeat_dim,\n        output_dim=num_hidden,\n        adj=adj,\n        input=infeats,\n        norm=norm,\n        activation=relay.nn.relu,\n    )\n)\nlayers.append(\n    GraphConv(\n        layer_name=\"layers.1\",\n        input_dim=num_hidden,\n        output_dim=num_classes,\n        adj=adj,\n        input=layers[-1],\n        norm=norm,\n        activation=None,\n    )\n)\n\n# Analyze free variables and generate Relay function\noutput = layers[-1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compile and run with TVM\n------------------------\n\nExport the weigths from PyTorch model to Python Dict\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model_params = {}\nfor param_tensor in torch_model.state_dict():\n    model_params[param_tensor] = torch_model.state_dict()[param_tensor].numpy()\n\nfor i in range(num_layers + 1):\n    params[\"layers.%d.weight\" % (i)] = model_params[\"layers.%d.weight\" % (i)]\n    params[\"layers.%d.bias\" % (i)] = model_params[\"layers.%d.bias\" % (i)]\n\n# Set the TVM build target\ntarget = \"llvm\"  # Currently only support `llvm` as target\n\nfunc = relay.Function(relay.analysis.free_vars(output), output)\nfunc = relay.build_module.bind_params_by_name(func, params)\nmod = tvm.IRModule()\nmod[\"main\"] = func\n# Build with Relay\nwith tvm.transform.PassContext(opt_level=0):  # Currently only support opt_level=0\n    lib = relay.build(mod, target, params=params)\n\n# Generate graph executor\ndev = tvm.device(target, 0)\nm = graph_executor.GraphModule(lib[\"default\"](dev))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Run the TVM model, test for accuracy and verify with DGL\n--------------------------------------------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "m.run()\nlogits_tvm = m.get_output(0).numpy()\nprint(\"Print the first five outputs from TVM execution\\n\", logits_tvm[:5])\n\nlabels = data.labels\ntest_mask = data.test_mask\n\nacc = evaluate(data, logits_tvm)\nprint(\"Test accuracy of TVM results: {:.2%}\".format(acc))\n\nimport tvm.testing\n\n# Verify the results with the DGL model\ntvm.testing.assert_allclose(logits_torch, logits_tvm, atol=1e-3)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}